#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：V2 
@File    ：T6Gurobi.py
@IDE     ：PyCharm 
@Author  ：郭星
@Date    ：2025/9/15 21:46
求解器求解
'''
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
求解器球技巧
"""
"""
用法 & 可扩展

替换实例：把 tiny_demo_instance() 改成你的数据读入（例如从 JSON/CSV 构造 FJSPInstance）。

求解参数：timelimit 控制时限（秒），mipgap 控制 MIP GAP；可在 build_fjsp_model 中添加更多参数（Heuristics, Cuts 等）。

收敛曲线：基于回调记录 Incumbent 与 Best Bound，求解结束后自动绘制。

复用：把 build_fjsp_model、solve_fjsp_with_gurobi 复制到你的工程中即可作为模块调用。

需要的话，我也可以参照这套写法，给你一版 混合流水车间（HFS） 的 Gurobi 模型（多阶段、每阶段并行机、目标最小化 makespan），同样带进度曲线与甘特图。
Gurobi MILP for Flexible Job Shop Scheduling (FJSP)
- x[j,k,m] ∈ {0,1}: operation (j,k) chooses machine m
- s[j,k]  ≥ 0: start time
- z[o1,o2,m] ∈ {0,1}: on machine m, o1 before o2 (pairwise sequencing)
- Precedence: within each job, op k+1 starts after op k completes
- Disjunctive: any two ops on the same machine do not overlap (big-M with assignment)
- Objective: minimize Cmax
- Callback: record incumbent & best bound over time to plot "MIP convergence"
- Gantt chart for final schedule

Usage:
    pip install gurobipy numpy matplotlib
    python gurobi_fjsp.py
"""

from dataclasses import dataclass
from typing import Dict, List, Tuple
import itertools
import numpy as np
import matplotlib.pyplot as plt

try:
    import gurobipy as gp
    from gurobipy import GRB
except Exception as e:
    raise SystemExit("请先安装并激活 Gurobi（pip install gurobipy 且配置许可证）。") from e


# ----------------------------
# 数据结构
# ----------------------------

@dataclass
class Operation:
    # 可行机器及加工时间（机器编号 -> 加工时间）
    machine_time: Dict[int, float]

@dataclass
class Job:
    ops: List[Operation]

@dataclass
class FJSPInstance:
    jobs: List[Job]
    m: int  # 机器总数

    @property
    def n_jobs(self) -> int:
        return len(self.jobs)

    @property
    def ops_list(self) -> List[Tuple[int, int]]:
        lst = []
        for j, job in enumerate(self.jobs):
            for k in range(len(job.ops)):
                lst.append((j, k))
        return lst

    @property
    def n_ops(self) -> int:
        return sum(len(job.ops) for job in self.jobs)

    def eligible(self, j: int, k: int) -> List[int]:
        return list(self.jobs[j].ops[k].machine_time.keys())

    def p(self, j: int, k: int, m: int) -> float:
        return float(self.jobs[j].ops[k].machine_time[m])

    def bigM(self) -> float:
        # 粗安全上界：各工序在可行机器上的最大加工时间总和 + 1
        return sum(max(op.machine_time.values()) for job in self.jobs for op in job.ops) + 1.0


# ----------------------------
# 建模
# ----------------------------

def build_fjsp_model(ins: FJSPInstance, timelimit: float = None, mipgap: float = 1e-4):
    """
    返回: model, x, s, z, Cmax（变量为 dict 索引）
    """
    mdl = gp.Model("FJSP_Gurobi")

    J = ins.n_jobs
    M = ins.m
    OPS = ins.ops_list
    BIGM = ins.bigM()

    # 决策变量
    # x[j,k,m]：分配变量
    x = {}
    for (j, k) in OPS:
        for m in ins.eligible(j, k):
            x[(j, k, m)] = mdl.addVar(vtype=GRB.BINARY, name=f"x_{j}_{k}_{m}")

    # s[j,k]：开始时间
    s = { (j, k): mdl.addVar(lb=0.0, vtype=GRB.CONTINUOUS, name=f"s_{j}_{k}") for (j,k) in OPS }

    # Cmax
    Cmax = mdl.addVar(lb=0.0, vtype=GRB.CONTINUOUS, name="Cmax")

    # 机器序变量 z
    z = {}
    for m in range(M):
        ops_m = [(j,k) for (j,k) in OPS if m in ins.eligible(j,k)]
        for (j1,k1), (j2,k2) in itertools.combinations(ops_m, 2):
            z[(j1,k1,j2,k2,m)] = mdl.addVar(vtype=GRB.BINARY, name=f"z_{j1}_{k1}_{j2}_{k2}_m{m}")

    mdl.update()

    # 目标
    mdl.setObjective(Cmax, GRB.MINIMIZE)

    # 约束
    # 1) 每个工序必须恰选一个机器
    for (j,k) in OPS:
        mdl.addConstr(gp.quicksum(x[(j,k,m)] for m in ins.eligible(j,k)) == 1, name=f"assign_{j}_{k}")

    # 2) 工序优先：s_{j,k+1} ≥ s_{j,k} + Σ p_{j,k,m} x_{j,k,m}
    for j in range(J):
        for k in range(len(ins.jobs[j].ops)-1):
            mdl.addConstr(
                s[(j,k+1)] >= s[(j,k)] + gp.quicksum(ins.p(j,k,m) * x[(j,k,m)] for m in ins.eligible(j,k)),
                name=f"prec_{j}_{k}_to_{k+1}"
            )

    # 3) 同机不重叠：二择序 + big-M + 仅在同机分配时激活
    for m in range(M):
        ops_m = [(j,k) for (j,k) in OPS if m in ins.eligible(j,k)]
        for (j1,k1), (j2,k2) in itertools.combinations(ops_m, 2):
            zvar = z[(j1,k1,j2,k2,m)]
            # o1 before o2
            mdl.addConstr(
                s[(j1,k1)] + ins.p(j1,k1,m)
                <= s[(j2,k2)] + BIGM*(1 - zvar) + BIGM*(2 - x[(j1,k1,m)] - x[(j2,k2,m)]),
                name=f"seq1_{j1}_{k1}_{j2}_{k2}_m{m}"
            )
            # o2 before o1
            mdl.addConstr(
                s[(j2,k2)] + ins.p(j2,k2,m)
                <= s[(j1,k1)] + BIGM*(zvar) + BIGM*(2 - x[(j1,k1,m)] - x[(j2,k2,m)]),
                name=f"seq2_{j1}_{k1}_{j2}_{k2}_m{m}"
            )

    # 4) Cmax 上界
    for (j,k) in OPS:
        mdl.addConstr(
            s[(j,k)] + gp.quicksum(ins.p(j,k,m) * x[(j,k,m)] for m in ins.eligible(j,k)) <= Cmax,
            name=f"cmax_{j}_{k}"
        )

    # 可选：收紧 s 的上界（帮助 M 更紧）
    for (j,k) in OPS:
        mdl.addConstr(s[(j,k)] <= BIGM, name=f"ubS_{j}_{k}")

    # 参数
    if timelimit is not None:
        mdl.Params.TimeLimit = float(timelimit)
    if mipgap is not None:
        mdl.Params.MIPGap = float(mipgap)
    mdl.Params.OutputFlag = 1

    return mdl, x, s, z, Cmax


# ----------------------------
# 进度回调（收敛曲线）
# ----------------------------

class MIPProgress:
    def __init__(self):
        self.times = []
        self.inc   = []   # incumbent
        self.bound = []   # best bound

    def callback(self, model, where):
        if where == GRB.Callback.MIP:
            t = model.cbGet(GRB.Callback.RUNTIME)
            best_inc = model.cbGet(GRB.Callback.MIP_OBJBST)
            best_bd  = model.cbGet(GRB.Callback.MIP_OBJBND)
            self.times.append(t)
            self.inc.append(best_inc)
            self.bound.append(best_bd)
        elif where == GRB.Callback.MIPSOL:
            t = model.cbGet(GRB.Callback.RUNTIME)
            best_inc = model.cbGet(GRB.Callback.MIPSOL_OBJ)
            # 在有新解时也记录一下（下界用最近一次 MIP 的）
            self.times.append(t)
            self.inc.append(best_inc)
            try:
                best_bd  = model.cbGet(GRB.Callback.MIP_OBJBND)
            except gp.GurobiError:
                best_bd = float('nan')
            self.bound.append(best_bd)

def plot_mip_convergence(times, incumbents, bounds, title="MIP Convergence"):
    if not times:  # 无回调数据
        return
    import numpy as np
    t = np.array(times, dtype=float)
    inc = np.array(incumbents, dtype=float)
    bd  = np.array(bounds, dtype=float)
    # 过滤无穷
    mask = np.isfinite(inc)
    t_inc, inc = t[mask], inc[mask]
    maskb = np.isfinite(bd)
    t_bd, bd = t[maskb], bd[maskb]

    plt.figure(figsize=(7,4))
    if len(t_bd) > 0:
        plt.step(t_bd, bd, where="post", label="Best Bound")
    if len(t_inc) > 0:
        plt.step(t_inc, inc, where="post", label="Best Incumbent")
    plt.xlabel("Time (s)")
    plt.ylabel("Objective (Cmax)")
    plt.title(title)
    plt.grid(True, alpha=0.3)
    plt.legend()
    plt.tight_layout()
    plt.show()


# ----------------------------
# 解读解并绘制甘特
# ----------------------------

def extract_schedule(ins: FJSPInstance, x, s) -> Tuple[float, Dict[Tuple[int,int], Tuple[int, float, float]]]:
    assign = {}
    makespan = 0.0
    for j in range(ins.n_jobs):
        for k, op in enumerate(ins.jobs[j].ops):
            chosen_m = None
            for m in op.machine_time.keys():
                if x[(j,k,m)].X > 0.5:
                    chosen_m = m
                    break
            if chosen_m is None:
                continue
            st = s[(j,k)].X
            en = st + ins.p(j,k,chosen_m)
            assign[(j,k)] = (chosen_m, st, en)
            makespan = max(makespan, en)
    return makespan, assign

def plot_gantt_fjsp(ins: FJSPInstance, assign: Dict[Tuple[int,int], Tuple[int, float, float]], title: str):
    bars = {m: [] for m in range(ins.m)}
    for (j,k), (m, st, en) in assign.items():
        bars[m].append((st, en-st, f"J{j}-O{k}"))
    fig, ax = plt.subplots(figsize=(12, 0.7*ins.m + 3))
    yticks, ylabels = [], []
    for row, m in enumerate(sorted(bars)):
        items = sorted(bars[m], key=lambda x: x[0])
        for (start, dur, label) in items:
            ax.barh(row, dur, left=start)
            ax.text(start + 0.5*dur, row, label, ha="center", va="center", fontsize=8)
        yticks.append(row); ylabels.append(f"M{m}")
    ax.set_yticks(yticks); ax.set_yticklabels(ylabels)
    ax.set_xlabel("Time"); ax.set_title(title)
    ax.grid(axis="x", alpha=0.3)
    plt.tight_layout()
    plt.show()


# ----------------------------
# 示例实例
# ----------------------------

def tiny_demo_instance() -> FJSPInstance:
    """
    小规模 FJSP 用于快速测试。
    5 台机器，4 个工件（部分工件 3~4 道工序），每道工序给出可行机与加工时间。
    """
    jobs: List[Job] = []

    # Job 0: 3 ops
    jobs.append(Job(ops=[
        Operation({0:4, 1:6}),         # O0,0
        Operation({1:5, 3:3, 4:4}),    # O0,1
        Operation({0:6, 2:5})          # O0,2
    ]))

    # Job 1: 4 ops
    jobs.append(Job(ops=[
        Operation({2:3, 3:5}),
        Operation({0:4, 4:6}),
        Operation({1:3, 2:4, 3:5}),
        Operation({0:2, 2:3})
    ]))

    # Job 2: 3 ops
    jobs.append(Job(ops=[
        Operation({0:5, 1:4}),
        Operation({2:6, 4:4}),
        Operation({1:5, 3:4})
    ]))

    # Job 3: 3 ops
    jobs.append(Job(ops=[
        Operation({3:3, 4:5}),
        Operation({0:5, 2:4}),
        Operation({1:4, 4:3})
    ]))

    return FJSPInstance(jobs=jobs, m=5)


# ----------------------------
# 主流程（可复用为函数）
# ----------------------------

def solve_fjsp_with_gurobi(ins: FJSPInstance, timelimit: float = 60.0, mipgap: float = 1e-4,
                           show_convergence: bool = True, show_gantt: bool = True):
    mdl, x, s, z, Cmax = build_fjsp_model(ins, timelimit, mipgap)

    progress = MIPProgress()
    print("Solving with Gurobi...")
    mdl.optimize(progress.callback)

    if mdl.SolCount == 0:
        print("未找到可行解（可能达到时限且无 incumbent）。")
        return None

    mk, assign = extract_schedule(ins, x, s)
    print(f"\nBest found makespan (Cmax): {mk:.2f}")
    for (j,k) in sorted(assign.keys()):
        m, st, en = assign[(j,k)]
        print(f"Job {j} Op {k} -> M{m}  start={st:.1f}  end={en:.1f}")

    if show_convergence:
        plot_mip_convergence(progress.times, progress.inc, progress.bound,
                             title="FJSP – Gurobi MIP Convergence")

    if show_gantt:
        plot_gantt_fjsp(ins, assign, title=f"FJSP – Gurobi schedule (Cmax={mk:.1f})")

    return {
        "model": mdl,
        "Cmax": mk,
        "assign": assign,
        "progress": progress
    }


# ----------------------------
# 脚本入口
# ----------------------------

def main():
    np.random.seed(0)
    ins = tiny_demo_instance()
    solve_fjsp_with_gurobi(ins, timelimit=60.0, mipgap=1e-4,
                           show_convergence=True, show_gantt=True)

if __name__ == "__main__":
    main()
